import numpy as np

def PSO(fitness, N, c1, c2, w, M, D):
    # % 待优化的目标函数：fitness
    # % 粒子数目： N
    # % 学习因子1：c1
    # % 学习银子2：c2
    # % 惯性权重：w
    # % 最大迭代次数：M
    # % 自变量的个数：D
    # % 目标函数取最小值时的自变量值：xm
    # % 目标函数的最小值：fv

    # 初始化
    x = np.zeros((N, D))
    v = np.zeros((N, D))
    for i in range(N):
        for j in range(D):
            x[i, j] = np.random.randn()
            v[i, j] = np.random.randn()
    
    p = np.zeros(N)
    y = np.zeros((N, D))
    for i in range(N):
        p[i] = fitness(x[i,:]) # 上一代适应度值
        y[i, :] = x[i,:]

    pg = x[N-1,:]
    for i in range(N-1):
        if fitness(x[i,:]) < fitness(pg):
            pg = x[i,:]

    # 迭代寻优
    for t in range(M):
        for i in range(N):
            v[i,:] = w*v[i,:] + c1*np.random.rand()*(y[i,:]-x[i,:]) + c2*np.random.rand()*(pg-x[i,:])
            x[i,:] = x[i,:] + v[i,:]
            if fitness(x[i,:]) < p[i]:
                p[i] = fitness(x[i,:])
                y[i,:] = x[i,:]
                if p[i] < fitness(pg):
                    pg = y[i,:]

    return pg, fitness(pg)

# 待求解函数
def fitness(x):
    F = 0
    for i in range(30):
        F = F + x[i]**2

    return F

# 求解
xm, fv = PSO(fitness, 40, 2, 2, 0.5, 100, 30)

print("最优解", fv)